﻿// DlgProcess.cpp: 实现文件
//

#include "pch.h"
#include "PE_Studio.h"
#include "afxdialogex.h"
#include "DlgProcess.h"
#include "DlgModule.h"
#include <TlHelp32.h>
#include <Psapi.h>
#include <strsafe.h>

// DlgProcess 对话框
IMPLEMENT_DYNAMIC(DlgProcess, CDialogEx)
DlgProcess::DlgProcess(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DLG_PROCESS, pParent)
{

}

DlgProcess::~DlgProcess()
{
}

void DlgProcess::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_LIST_PROCESS, m_ProcessListCtrl);
}


// 消息映射
BEGIN_MESSAGE_MAP(DlgProcess, CDialogEx)
	ON_NOTIFY(NM_RCLICK, IDC_LIST_PROCESS, &DlgProcess::OnNMRClickListProcess)
	ON_COMMAND(ID_40001, &DlgProcess::On40001)
	ON_COMMAND(ID_40002, &DlgProcess::On40002)
	ON_COMMAND(ID_40003, &DlgProcess::On40003)
	ON_COMMAND(ID_40004, &DlgProcess::On40004)
	ON_COMMAND(ID_40005, &DlgProcess::On40005)
END_MESSAGE_MAP()


// DlgProcess 消息处理程序
BOOL DlgProcess::OnInitDialog()
{
	CDialogEx::OnInitDialog();

	// TODO:  在此添加额外的初始化
	m_ProcessListCtrl.SetListCtrl(TRUE, TRUE, TRUE, FALSE, TRUE);		// 禁止排序
    m_ProcessListCtrl.InsertColumn(0, TEXT(""), LVCFMT_CENTER, 50);
    m_ProcessListCtrl.InsertColumn(1, TEXT("进程名"), LVCFMT_CENTER, 100);
    m_ProcessListCtrl.InsertColumn(2, TEXT("父进程"), LVCFMT_CENTER, 60);
    m_ProcessListCtrl.InsertColumn(3, TEXT("进程ID"), LVCFMT_CENTER, 60);
    m_ProcessListCtrl.InsertColumn(4, TEXT("线程数"), LVCFMT_CENTER, 60);
    m_ProcessListCtrl.InsertColumn(5, TEXT("优先级"), LVCFMT_CENTER, 60);
    m_ProcessListCtrl.InsertColumn(6, TEXT("进程路径"), LVCFMT_CENTER, 230);
    m_ProcessListCtrl.DeleteColumn(0);

	if (!InitProcessList())
	{
		return FALSE;
	}

	return TRUE;  // return TRUE unless you set the focus to a control
	// 异常: OCX 属性页应返回 FALSE
}


// 右键菜单
void DlgProcess::OnNMRClickListProcess(NMHDR* pNMHDR, LRESULT* pResult)
{
	LPNMITEMACTIVATE pNMItemActivate = reinterpret_cast<LPNMITEMACTIVATE>(pNMHDR);
	// TODO: 在此添加控件通知处理程序代码

	POINT pt = {};
	GetCursorPos(&pt);
	CMenu popMenu;
	CMenu* pPopMenu = NULL;
	popMenu.LoadMenu(IDR_MENU2);
	pPopMenu = popMenu.GetSubMenu(0);
	pPopMenu->TrackPopupMenu(TPM_LEFTALIGN, pt.x, pt.y, this, NULL);

	*pResult = 0;
}

// 初始化进程列表
DWORD WINAPI DlgProcess::ThreadProc(LPVOID lpParam)
{
	DlgProcess* This = (DlgProcess*)lpParam;
	HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (INVALID_HANDLE_VALUE == hProcessSnap)
	{
		return 0;
	}
	/*
		typedef struct tagPROCESSENTRY32 {
		  DWORD     dwSize;
		  DWORD     cntUsage;
		  DWORD     th32ProcessID;			//进程标识符。
		  ULONG_PTR th32DefaultHeapID;
		  DWORD     th32ModuleID;
		  DWORD     cntThreads;				//线程数。
		  DWORD     th32ParentProcessID;	//创建此进程的进程的标识符（其父进程）
		  LONG      pcPriClassBase;			//线程的基本优先级。
		  DWORD     dwFlags;
		  CHAR      szExeFile[MAX_PATH];	//进程的可执行文件的名称。
		} PROCESSENTRY32;
	*/
	PROCESSENTRY32 pe32 = { 0 };
	pe32.dwSize = sizeof(PROCESSENTRY32);
	TCHAR cParentProcessID[0x10] = { 0 };
	TCHAR cProcessID[0x10] = { 0 };
	TCHAR ccntThreds[0x10] = { 0 };
	TCHAR cpcPriClassBase[0x10] = { 0 };
	TCHAR cProcessPath[MAX_PATH] = { 0 };
	UINT uIndex = 0;
	if (!Process32First(hProcessSnap, &pe32))
	{
		::CloseHandle(hProcessSnap);
		return 0;
	}
	do
	{
		HANDLE hProcess = ::OpenProcess(PROCESS_ALL_ACCESS, FALSE, pe32.th32ProcessID);
		::GetModuleFileNameEx(hProcess, NULL, cProcessPath, MAX_PATH);
		if (NULL == hProcess)
		{
			StringCchPrintf(cProcessPath, MAX_PATH, TEXT("System Module"));
		}

		// 添加进程进列表
		StringCchPrintf(cParentProcessID, 0x10, TEXT("%d"), pe32.th32ParentProcessID);
		StringCchPrintf(cProcessID, 0x10, TEXT("%d"), pe32.th32ProcessID);
		StringCchPrintf(ccntThreds, 0x10, TEXT("%d"), pe32.cntThreads);
		StringCchPrintf(cpcPriClassBase, 0x10, TEXT("%d"), pe32.pcPriClassBase);

		This->m_ProcessListCtrl.InsertItem(uIndex, pe32.szExeFile);
		This->m_ProcessListCtrl.SetItemText(uIndex, 1, cParentProcessID);
		This->m_ProcessListCtrl.SetItemText(uIndex, 2, cProcessID);
		This->m_ProcessListCtrl.SetItemText(uIndex, 3, ccntThreds);
		This->m_ProcessListCtrl.SetItemText(uIndex, 4, cpcPriClassBase);
		This->m_ProcessListCtrl.SetItemText(uIndex, 5, cProcessPath);

		uIndex++;

	} while (Process32Next(hProcessSnap, &pe32));

	::CloseHandle(hProcessSnap);
    return 0;
}
BOOL DlgProcess::InitProcessList()
{
	HANDLE hThread = ::CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)ThreadProc, this, 0, NULL);
	if (hThread != NULL)
	{
		::CloseHandle(hThread);
	}
	return 0;
}

// 设置访问令牌中的权限
BOOL DlgProcess::EnableDebugPrivilege(BOOL bEnable)
{
    BOOL bRet = FALSE;
    HANDLE hToken = NULL;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken))
	{
        TOKEN_PRIVILEGES tp;
        tp.PrivilegeCount = 1;												// 给PrivilegeCount赋值
        LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tp.Privileges[0].Luid);	// 给Luid赋值
        tp.Privileges[0].Attributes = bEnable ? SE_PRIVILEGE_ENABLED : 0;	// 给Attributes赋值

        AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), NULL, NULL);
		bRet = (GetLastError() == ERROR_SUCCESS);
        CloseHandle(hToken);
    }
    return bRet;
}

// 创建进程
void DlgProcess::On40001()
{
	// TODO: 在此添加命令处理程序代码

	CFileDialog fileDlg(TRUE, TEXT("exe"), TEXT("*.exe"), OFN_FILEMUSTEXIST, TEXT("Executable Files(*.exe)|*.exe|All Files(*.*)|*.*||"));
	if (fileDlg.DoModal() == IDOK)
	{
        CString csProcessPathName = fileDlg.GetPathName();
		STARTUPINFO si = { 0 };
		PROCESS_INFORMATION pi = { 0 };
		si.cb = sizeof(si);

		if (::CreateProcess(csProcessPathName, \
			NULL, NULL, NULL, FALSE, \
			CREATE_NEW_CONSOLE, \
			NULL, NULL, \
			& si, &pi))
		{
            ::CloseHandle(pi.hThread);     //结束新创建主线程句柄
            ::CloseHandle(pi.hProcess);    //结束新创建主进程句柄
		}
	}
}

// 结束进程
void DlgProcess::On40002()
{
	// TODO: 在此添加命令处理程序代码
	EnableDebugPrivilege(TRUE);	// 提升权限
	// 获取m_ProcessListCtrl选中行
	DWORD dwPosition = (DWORD)m_ProcessListCtrl.GetFirstSelectedItemPosition();
	dwPosition -= 1;
	CString csProcessID = m_ProcessListCtrl.GetItemText(dwPosition, 2);
	DWORD dwProcessID = _tcstoul(csProcessID, NULL, 10);
	HANDLE hProcess = ::OpenProcess(PROCESS_TERMINATE, FALSE, dwProcessID);
	if (NULL == hProcess)
	{
		AfxMessageBox(TEXT("OpenProcess Failed!"));
		return;
	}
	// 用API结束 进程 
	if (::TerminateProcess(hProcess, 0))
	{
		AfxMessageBox(TEXT("TerminateProcess Success!"));
	}
	else
	{
		AfxMessageBox(TEXT("TerminateProcess Failed!"));
		return;
	}
	EnableDebugPrivilege(FALSE);	// 恢复权限

	m_ProcessListCtrl.DeleteAllItems();
	InitProcessList();
}

// 刷新列表
void DlgProcess::On40003()
{
    m_ProcessListCtrl.DeleteAllItems();
    InitProcessList();
}

// 查看模块
void DlgProcess::On40004()
{
	DlgModule m_DialogModule(this);
	m_DialogModule.DoModal();
}

// 打开文件所在文件夹
void DlgProcess::On40005()
{
	int nIndex = -1;
	CString szExePath = TEXT("");
	nIndex = m_ProcessListCtrl.GetSelectionMark();
	if (nIndex < 0)
    {
		MessageBox(TEXT("打开所在文件夹 - 请先选中进程"), TEXT("提醒"), MB_OK | MB_ICONINFORMATION);
        return;
    }
	szExePath = m_ProcessListCtrl.GetItemText(nIndex, 5);				// 获取路径
    if (PathFileExists(szExePath))
    {
        TCHAR szParams[MAX_PATH] = { 0 };
        _stprintf_s(szParams, TEXT("/select, %s"), szExePath);
		ShellExecute(NULL, NULL, TEXT("explorer.exe"), szParams, NULL, SW_NORMAL);	// 打开文件所在目录
        return;
    }
}

